package bm.com.ai.server.rag.service;

import bm.com.ai.server.dto.MessageResponse;
import bm.com.ai.server.rag.config.RagConfig;
import bm.com.ai.server.vector.service.QdrantVectorService;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.huggingface.HuggingFaceEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class RagService {

    @Resource
    private final RagConfig ragConfig;

    @Resource
    private final QdrantVectorService vectorStore;

    @Resource
    private EmbeddingModel embeddingModel;

    private final ThreadLocal<List<MessageResponse.RagSource>> lastSources = new ThreadLocal<>();

    /**
     * 检索相关上下文
     */
    public String retrieveRelevantContext(String query) {
        if (!ragConfig.isEnabled()) {
            return "";
        }

        try {

            // 2. 在向量存储中查找相关文档
            List<EmbeddingMatch<TextSegment>> embeddingMatch = vectorStore.getEmbeddingMatch(query);

//            // 3. 保存来源信息供后续使用
//            saveSources(embeddingMatch);

            // 4. 构建上下文字符串
            return buildContextString(embeddingMatch);

        } catch (Exception e) {
            log.error("Error retrieving relevant context for query: {}", query, e);
            // RAG检索失败不应导致整个请求失败，返回空上下文
            return "";
        }
    }

    /**
     * 构建上下文字符串
     */
    private String buildContextString(List<EmbeddingMatch<TextSegment>> matches) {
        if (matches.isEmpty()) {
            return "";
        }

        StringBuilder context = new StringBuilder();
        context.append("相关信息：\n");

        for (int i = 0; i < matches.size(); i++) {
            EmbeddingMatch<TextSegment> match = matches.get(i);
            context.append(i + 1).append(". ").append(match.embedded().text()).append("\n");
        }

        return context.toString();
    }

    /**
     * 保存来源信息
     */
    private void saveSources(List<EmbeddingMatch<TextSegment>> matches) {
        List<MessageResponse.RagSource> sources = new ArrayList<>();

        for (EmbeddingMatch<TextSegment> match : matches) {
            MessageResponse.RagSource source = MessageResponse.RagSource.builder()
                    .id(match.embedded().text())
                    .content(match.embedded().text())
                    .similarityScore(match.score())
                    .source(String.valueOf(match.embedded().metadata().containsKey("source")))
                    .metadata((Map<String, String>) match.embedded().metadata())
                    .build();
            sources.add(source);
        }

        lastSources.set(sources);
    }

    /**
     * 获取上次检索的来源信息
     */
    public List<MessageResponse.RagSource> getLastSources() {
        List<MessageResponse.RagSource> sources = lastSources.get();
        return sources != null ? new ArrayList<>(sources) : new ArrayList<>();
    }

    /**
     * 向量存储添加文档
     */
    public void addDocuments(List<TextSegment> segments) {
        if (!ragConfig.isEnabled() || segments.isEmpty()) {
            return;
        }

        try {
            // 生成嵌入
            List<Embedding> embeddings = embeddingModel.embedAll(segments)
                    .content();

            // 添加到向量存储
            vectorStore.addAll(embeddings, segments);

            log.info("Added {} documents to vector store", segments.size());
        } catch (Exception e) {
            log.error("Error adding documents to vector store", e);
            throw new RuntimeException("Error adding documents to vector store", e);
        }
    }

    /**
     * 向向量存储添加单个文档
     */
    public void addDocument(TextSegment segment) {
        if (!ragConfig.isEnabled()) {
            return;
        }

        try {
            // 生成嵌入
            Embedding embedding = embeddingModel.embed(segment.text()).content();

            // 添加到向量存储
            vectorStore.add(embedding, segment);

            log.debug("Added document {} to vector store", segment.toString());
        } catch (Exception e) {
            log.error("Error adding document to vector store", e);
            throw new RuntimeException("Error adding document to vector store", e);
        }
    }
}
